"""
使用 Gradio 启动一个基于手写数字的 KNN 预测 Web 应用。

依赖：gradio, numpy, pillow(PIL), opencv-python（可选）
模型文件：与此脚本同目录下的 best_knn.pkl（由 optimal_knn.py 生成）
"""

# TODO: 导入必要的库和模块
from __future__ import annotations

import io
import pickle
from pathlib import Path
from typing import Tuple, Any

import numpy as np
from PIL import Image, ImageOps, ImageFilter

import gradio as gr
try:
	import cv2  # 可选：用于更快的 Otsu 与膨胀
except Exception:
	cv2 = None


# TODO: 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
MODEL_PATH = Path(__file__).resolve().parent / "best_knn.pkl"
_MODEL = None
_BEST_K = None
_BEST_ACC = None

if MODEL_PATH.exists():
	with MODEL_PATH.open("rb") as f:
		payload = pickle.load(f)
		_MODEL = payload.get("model")
		_BEST_K = payload.get("best_k")
		_BEST_ACC = payload.get("best_acc")
else:
	# 延迟报错到预测阶段，使应用能显示清晰提示
	pass


def _to_pil(input_obj: Any) -> Image.Image:
	"""将 Sketchpad 可能返回的多种类型(dict/ndarray/PIL)统一转为 PIL.Image.

	兼容格式：
	- PIL.Image
	- numpy.ndarray (H,W) 或 (H,W,3|4)
	- dict：常见键 image/composite/background/layers（取首个可解析的图层）
	"""
	if input_obj is None:
		raise ValueError("未检测到图像")

	# 已是 PIL
	if isinstance(input_obj, Image.Image):
		pil = input_obj
	# numpy 数组
	elif isinstance(input_obj, np.ndarray):
		pil = Image.fromarray(input_obj)
	# gradio 某些版本的 Sketchpad/ImageEditor 可能返回 dict
	elif isinstance(input_obj, dict):
		candidate = None
		for key in ("image", "composite", "background"):
			if key in input_obj and input_obj[key] is not None:
				candidate = input_obj[key]
				break
		# 尝试 layers（若存在），取最后一层
		if candidate is None and isinstance(input_obj.get("layers"), list) and input_obj["layers"]:
			last = input_obj["layers"][-1]
			if isinstance(last, dict):
				candidate = last.get("image") or last.get("composite")
			else:
				candidate = last

		if candidate is None:
			raise ValueError("无法从输入字典解析图像内容")

		if isinstance(candidate, Image.Image):
			pil = candidate
		elif isinstance(candidate, np.ndarray):
			pil = Image.fromarray(candidate)
		else:
			# 有些实现可能是 base64/bytes，这里先简单兜底
			raise ValueError("不支持的图像数据类型: " + type(candidate).__name__)
	else:
		raise ValueError("不支持的输入类型: " + type(input_obj).__name__)

	# 去除 alpha：合成到白底，避免反相后出现暗底
	if pil.mode in ("RGBA", "LA"):
		bg = Image.new("RGBA", pil.size, (255, 255, 255, 255))
		pil = Image.alpha_composite(bg, pil.convert("RGBA")).convert("RGB")

	return pil


def _otsu_threshold(arr: np.ndarray) -> int:
	"""简易 Otsu 自动阈值计算，返回 0..255 的整型阈值。"""
	# 优先使用 OpenCV 的 Otsu（C 实现更快）
	if cv2 is not None:
		try:
			_arr = arr.astype(np.uint8, copy=False)
			_, th = cv2.threshold(_arr, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
			return int(_)
		except Exception:
			pass
	hist, _ = np.histogram(arr, bins=256, range=(0, 255))
	total = arr.size
	sum_total = np.dot(np.arange(256), hist)
	sumB = 0.0
	wB = 0.0
	max_var = -1.0
	thresh = 0
	for t in range(256):
		wB += hist[t]
		if wB == 0:
			continue
		wF = total - wB
		if wF == 0:
			break
		sumB += t * hist[t]
		mB = sumB / wB
		mF = (sum_total - sumB) / wF
		varBetween = wB * wF * (mB - mF) ** 2
		if varBetween > max_var:
			max_var = varBetween
			thresh = t
	return int(thresh)


def _preprocess_to_digits(
	img: Image.Image,
	invert: bool = True,
	thresh: int = 10,
	padding: int = 4,
	blur_radius: float = 1.0,
	dilate: int = 0,
	auto_thresh: bool = False,
) -> np.ndarray:
	"""
	将任意输入图像预处理为 sklearn digits (8x8, 值域 0..16) 的扁平 1x64 向量。

	步骤：转灰度 -> 反相(使数字为亮、背景暗) -> 缩放到 8x8 -> 归一化到 0..16 -> flatten
	"""
	# 转灰度
	gray = img.convert("L")

	# 可选反相：sklearn digits 更像黑底白字
	if invert:
		gray = ImageOps.invert(gray)

	# 居中裁剪：找到笔迹的外接框，做少量 padding 后裁剪，提升缩放后有效占比
	arr_full = np.array(gray)
	# 自动阈值（可选）
	if auto_thresh:
		t = _otsu_threshold(arr_full)
	else:
		t = max(0, min(255, int(thresh)))
	# 简单二值化（可调阈值/自动阈值）
	mask = arr_full > t
	if mask.any():
		ys, xs = np.where(mask)
		y0, y1 = ys.min(), ys.max()
		x0, x1 = xs.min(), xs.max()
		pad = max(0, int(padding))
		y0 = max(0, y0 - pad); y1 = min(arr_full.shape[0] - 1, y1 + pad)
		x0 = max(0, x0 - pad); x1 = min(arr_full.shape[1] - 1, x1 + pad)
		arr_cropped = arr_full[y0:y1+1, x0:x1+1]
	else:
		# 空图
		arr_cropped = arr_full

	# 填充成正方形（黑底），保持纵横比
	h, w = arr_cropped.shape
	side = max(h, w)
	canvas = np.zeros((side, side), dtype=arr_cropped.dtype)
	y_off = (side - h) // 2
	x_off = (side - w) // 2
	canvas[y_off:y_off+h, x_off:x_off+w] = arr_cropped
	cropped = Image.fromarray(canvas)

	# 简易形态学：可选膨胀，增强细线条
	if dilate and dilate > 0:
		k = 2 * int(dilate) + 1
		if cv2 is not None:
			try:
				kernel = np.ones((k, k), np.uint8)
				cropped_np = np.array(cropped)
				dilated = cv2.dilate(cropped_np, kernel, iterations=1)
				cropped = Image.fromarray(dilated)
			except Exception:
				cropped = cropped.filter(ImageFilter.MaxFilter(size=k))
		else:
			# 使用最大滤波作为近似膨胀
			cropped = cropped.filter(ImageFilter.MaxFilter(size=k))

	# 轻度模糊，模拟 digits 的抗锯齿感（可调）
	if blur_radius and blur_radius > 0:
		cropped = cropped.filter(ImageFilter.GaussianBlur(radius=float(blur_radius)))

	# 调整大小到 8x8，适配不同 Pillow 版本的重采样常量
	try:
		resample = Image.Resampling.LANCZOS  # Pillow >= 9.1
	except Exception:
		# Pillow < 9.1 兼容
		resample = getattr(Image, "LANCZOS", Image.BICUBIC)
	small = cropped.resize((8, 8), resample)

	arr = np.array(small).astype(np.float32)  # 0..255
	# 归一化到 0..16
	arr = (arr / 255.0) * 16.0
	arr = np.clip(arr, 0.0, 16.0)
	# 扁平化为 (1, 64)
	flat = arr.flatten().reshape(1, -1)
	return flat


# TODO: 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(
	input_img: Any,
	invert: bool = True,
	thresh: int = 10,
	padding: int = 4,
	blur_radius: float = 1.0,
	dilate: int = 0,
	preview_size: int = 128,
	show_grid: bool = False,
	auto_thresh: bool = True,
	fast_mode: bool = False,
):
	"""接收来自 Sketchpad 的输入（可能是 PIL、ndarray 或 dict），返回预测类别、说明与预处理可视化。"""
	try:
		if _MODEL is None:
			return -1, "未找到模型文件 best_knn.pkl。请先运行 optimal_knn.py 训练并生成模型。"

		if input_img is None:
			return -1, "未检测到图像，请在画板中写一个数字后再提交。"

		pil_img = _to_pil(input_img)
		if fast_mode:
			# 极速路径：禁用较重步骤
			features = _preprocess_to_digits(
				pil_img,
				invert=invert,
				thresh=thresh,
				padding=padding,
				blur_radius=0.0,
				dilate=0,
				auto_thresh=False,
			)
		else:
			features = _preprocess_to_digits(
				pil_img,
				invert=invert,
				thresh=thresh,
				padding=padding,
				blur_radius=blur_radius,
				dilate=dilate,
				auto_thresh=auto_thresh,
			)
		pred = _MODEL.predict(features)
		label = int(pred[0])
		if _BEST_K is not None and _BEST_ACC is not None:
			info = f"预测: {label} (K={_BEST_K}, 验证准确率≈{_BEST_ACC:.4f})"
		else:
			info = f"预测: {label}"
		# 生成 8x8 预处理可视化图
		img8 = (features.reshape(8, 8) / 16.0 * 255.0).astype(np.uint8)
		pv = 128 if fast_mode else max(32, int(preview_size))
		viz = Image.fromarray(img8, mode="L").resize((pv,) * 2, Image.NEAREST)
		if show_grid:
			# 叠加网格，帮助观察每个像素格
			viz = viz.convert("RGB")
			import PIL.ImageDraw as ImageDraw
			draw = ImageDraw.Draw(viz)
			step = viz.size[0] // 8
			for i in range(1, 8):
				draw.line([(i * step, 0), (i * step, viz.size[1])], fill=(200, 200, 200), width=1)
				draw.line([(0, i * step), (viz.size[0], i * step)], fill=(200, 200, 200), width=1)
		return label, info, viz
	except Exception as e:
		# 捕获异常，避免 Gradio 输出“错误”红框；改为以文本形式提示
		return -1, f"发生错误：{e}", None


# TODO: 创建Gradio接口，这个接口将用于用户输入和显示预测结果
with gr.Blocks(title="KNN 手写数字识别") as demo:
	gr.Markdown("# KNN 手写数字识别\n在画板中写下一个数字并提交，模型将给出预测结果。")

	with gr.Row():
		sketch = gr.Sketchpad(
			label="画板(建议写粗一点)",
			type="numpy",  # 返回 numpy.ndarray，可减少返回 dict 的情况
			height=224,
			width=224,
		)
		with gr.Column():
			invert_chk = gr.Checkbox(value=True, label="反相(白字黑底)")
			with gr.Row():
				thresh_sld = gr.Slider(0, 255, value=25, step=1, label="阈值(找前景)")
				pad_sld = gr.Slider(0, 32, value=6, step=1, label="外接框边距")
			with gr.Row():
				blur_sld = gr.Slider(0, 3, value=0.5, step=0.5, label="模糊半径")
				dilate_sld = gr.Slider(0, 5, value=1, step=1, label="膨胀(加粗)")
			with gr.Row():
				prev_sld = gr.Slider(64, 256, value=192, step=32, label="预览尺寸")
				grid_chk = gr.Checkbox(value=False, label="显示网格")
			auto_thresh_chk = gr.Checkbox(value=True, label="自动阈值(Otsu)")
			fast_chk = gr.Checkbox(value=False, label="极速模式(更快，可能略降精度)")
			btn = gr.Button("提交")
			out_label = gr.Number(label="预测数字", precision=0)
			out_text = gr.Textbox(label="说明")
			out_preview = gr.Image(label="预处理8x8视图", type="pil")

			# 清除按钮：只清空画板与预览，不更改参数，也不改写预测结果
			try:
				gr.ClearButton(
					components=[sketch, out_preview],
					value="清除",
				)
			except Exception:
				clear_btn = gr.Button("清除")
				def _clear_only_canvas_preview():
					# 仅返回画板和预览的清空结果；其他输出/参数不变
					return None, None
				clear_btn.click(
					_clear_only_canvas_preview,
					inputs=None,
					outputs=[sketch, out_preview],
				)

	btn.click(
		predict_digit,
		inputs=[sketch, invert_chk, thresh_sld, pad_sld, blur_sld, dilate_sld, prev_sld, grid_chk, auto_thresh_chk, fast_chk],
		outputs=[out_label, out_text, out_preview],
	)


# TODO: 启动Gradio接口，用户可以通过这个接口进行交互
if __name__ == "__main__":
	# 在本地运行，无需创建公共链接；如需外网访问可设置 share=True
	demo.launch()
